#include "tunnel.h"
#include "router.h"
#include "server.h"
#include "../common/codec.h"
#include "../common/factory.h"



Tunnel::Tunnel (Server* server, Router* router, jlib::RSA& rsa, s32 socket, u32 ipv4, u16 port)
: _server(server)
, _router(router)
, _rsa(rsa)
, _socket(socket)
, _ipv4(ipv4)
, _port(port) {
  TRACE
}


Tunnel::~Tunnel () {
  TRACE
}


void Tunnel::startup () {
  TRACE
  _upWorker = new std::thread(&Tunnel::upLoop, this);
  _upWorker->detach();
}


void Tunnel::shutdown () {
  TRACE
  std::lock_guard<std::mutex> lock {_close_mutex};
  if (_closed)
    return;
  _closed = true;
  if (_upWorker)
    delete _upWorker;
  if (_downWorker)
    delete _downWorker;
  _router->releaseMapping(this, _ipv4, _port);
  _server->delTunnel(this);
  close(_socket);
  release();
}


void Tunnel::upLoop () {
  TRACE
  if (!handshake()) {
    log_e("Tunnel: handshake failed");
    return;
  }
  _downWorker = new std::thread(&Tunnel::downLoop, this);
  _downWorker->detach();

  u8 buf[IpPackage::MAX_SIZE + 16];
  while (!_closed) {
    auto nread = recv(_socket, buf, IpPackage::MAX_SIZE + 16, 0);
    if (nread <= 0) {
      log_e("Tunnel: read data error");
      shutdown();
      return;
    }
    _cache_mutex.lock();
    _cache.normalize();
    _cache.write(buf, (uwl)nread);
    _cache_mutex.unlock();
    Factory::executor().execute([this] () {
      for (;;) {
        _cache_mutex.lock();
        auto data = Codec::decode(_cache, _aes, _aes_key);
        _cache_mutex.unlock();
        if (!data.available())
          break;
        auto pkg = Codec::parse(data);
        if (!pkg)
          break;
        if (!pkg->err.empty()) {
          log_w("Tunnel.package: %s", pkg->err.c_str());
          continue;
        }
        _router->forwardUp(this, pkg);
      }
    });
  }
}


void Tunnel::downLoop () {
  TRACE
  while (!_closed) {
    auto pkg = _queue.poll();
    auto data = Codec::encode(pkg->buf, _aes, _aes_key);
    auto nwrite = send(_socket, data.readIter(), data.readable(), 0);
    if (nwrite <= 0) {
      log_e("Tunnel: write data error");
      shutdown();
      return;
    }
  }
}


bool Tunnel::handshake () {
  TRACE
  u8 buf[1024];
  abort_assert(!_cache.readable());
  while (!_closed) {
    if (_cache.readIndex() > 1024 * 1024)
      _cache.normalize();
    auto nread = read(_socket, buf, 1024);
    if (nread <= 0) {
      log_e("Tunnel: read key error");
      shutdown();
      return false;
    }
    _cache.write(buf, (uwl)nread);
    abort_assert(_cache.readable() == (uwl)nread);
    auto data = Codec::decode(_cache, _rsa);
    if (!data.available())
      continue;
    if (data.readable() != 32) {
      log_e("Tunnel: handshake error");
      return false;
    }
    memcpy(_aes_key, data.base(), 32);
    _aes.setup(_aes_key, 256);
    return true;
  }
  return false;
}


void Tunnel::write (Ptr<IpPackage> pkg) {
  TRACE
  std::lock_guard<std::mutex> lock {_close_mutex};
  if (_closed)
    return;
  _queue.push(std::move(pkg));
}